use ssh2::Session;
use tokio::net::{TcpListener as TokioTcpListener, TcpStream as TokioTcpStream};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use std::sync::Arc;
use tokio::sync::Mutex;
use anyhow::{Result, anyhow};
use uuid::Uuid;
use std::collections::HashMap;
use log::{info, error, debug};

/// 隧道类型
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub enum TunnelType {
    /// 本地端口转发 (Local Port Forwarding)
    /// 将本地端口转发到远程主机的端口
    Local {
        local_port: u16,
        remote_host: String,
        remote_port: u16,
    },
    /// 远程端口转发 (Remote Port Forwarding)
    /// 将远程端口转发到本地主机的端口
    Remote {
        remote_port: u16,
        local_host: String,
        local_port: u16,
    },
    /// 动态端口转发 (Dynamic Port Forwarding / SOCKS Proxy)
    /// 创建SOCKS代理
    Dynamic {
        local_port: u16,
    },
}

/// 隧道状态
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub enum TunnelStatus {
    Stopped,
    Starting,
    Running,
    Error(String),
}

/// 隧道配置
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TunnelConfig {
    pub id: String,
    pub name: String,
    pub tunnel_type: TunnelType,
    pub connection_id: String, // 关联的SSH连接ID
    pub auto_start: bool,
    pub description: Option<String>,
}

impl TunnelConfig {
    pub fn new(
        name: String,
        tunnel_type: TunnelType,
        connection_id: String,
    ) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            name,
            tunnel_type,
            connection_id,
            auto_start: false,
            description: None,
        }
    }
}

/// 隧道实例
pub struct Tunnel {
    pub config: TunnelConfig,
    pub status: Arc<Mutex<TunnelStatus>>,
    listener: Option<TokioTcpListener>,
    active_connections: Arc<Mutex<HashMap<String, TokioTcpStream>>>,
}

impl Tunnel {
    pub fn new(config: TunnelConfig) -> Self {
        Self {
            config,
            status: Arc::new(Mutex::new(TunnelStatus::Stopped)),
            listener: None,
            active_connections: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// 启动隧道
    pub async fn start(&mut self, ssh_session: Arc<Mutex<Session>>) -> Result<()> {
        info!("启动隧道: {} ({})", self.config.name, self.config.id);
        
        {
            let mut status = self.status.lock().await;
            *status = TunnelStatus::Starting;
        }

        match &self.config.tunnel_type {
            TunnelType::Local { local_port, remote_host, remote_port } => {
                self.start_local_forward(*local_port, remote_host.clone(), *remote_port, ssh_session).await?
            }
            TunnelType::Remote { remote_port, local_host, local_port } => {
                self.start_remote_forward(*remote_port, local_host.clone(), *local_port, ssh_session).await?
            }
            TunnelType::Dynamic { local_port } => {
                self.start_dynamic_forward(*local_port, ssh_session).await?
            }
        }

        {
            let mut status = self.status.lock().await;
            *status = TunnelStatus::Running;
        }

        info!("隧道启动成功: {}", self.config.name);
        Ok(())
    }

    /// 启动本地端口转发
    async fn start_local_forward(
        &mut self,
        local_port: u16,
        remote_host: String,
        remote_port: u16,
        ssh_session: Arc<Mutex<Session>>,
    ) -> Result<()> {
        let bind_addr = format!("127.0.0.1:{}", local_port);
        let listener = TokioTcpListener::bind(&bind_addr).await
            .map_err(|e| anyhow!("绑定本地端口失败 {}: {}", local_port, e))?;
        
        info!("本地端口转发监听: {} -> {}:{}", bind_addr, remote_host, remote_port);
        
        // 先存储 listener
        self.listener = Some(listener);
        
        // 获取 listener 的引用用于 spawn
        if let Some(ref listener) = self.listener {
            let active_connections = Arc::clone(&self.active_connections);
            let _tunnel_id = self.config.id.clone();
            
            // 克隆必要的数据用于 spawn
            let listener_addr = listener.local_addr().map_err(|e| anyhow!("获取监听地址失败: {}", e))?;
            
            // 重新绑定一个新的 listener 用于 spawn
            let spawn_listener = TokioTcpListener::bind(listener_addr).await
                .map_err(|e| anyhow!("重新绑定端口失败: {}", e))?;
            
            // 启动接受连接的任务
            tokio::spawn(async move {
                loop {
                    match spawn_listener.accept().await {
                        Ok((local_stream, addr)) => {
                            debug!("接受本地连接: {}", addr);
                            
                            let connection_id = Uuid::new_v4().to_string();
                            let ssh_session = Arc::clone(&ssh_session);
                            let remote_host = remote_host.clone();
                            let active_connections = Arc::clone(&active_connections);
                            let conn_id = connection_id.clone();
                            
                            // 为每个连接启动转发任务
                            tokio::spawn(async move {
                                if let Err(e) = Self::handle_local_forward_connection(
                                    local_stream,
                                    ssh_session,
                                    remote_host,
                                    remote_port,
                                ).await {
                                    error!("本地转发连接处理失败: {}", e);
                                }
                                
                                // 清理连接
                                let mut connections = active_connections.lock().await;
                                connections.remove(&conn_id);
                            });
                        }
                        Err(e) => {
                            error!("接受连接失败: {}", e);
                            break;
                        }
                    }
                }
            });
        }
        
        Ok(())
    }

    /// 处理本地转发连接
    async fn handle_local_forward_connection(
        mut local_stream: TokioTcpStream,
        ssh_session: Arc<Mutex<Session>>,
        remote_host: String,
        remote_port: u16,
    ) -> Result<()> {
        // 注意：这里需要使用ssh2的channel来建立到远程主机的连接
        // 由于ssh2库主要是同步的，这里需要在tokio的blocking任务中处理
        let (_local_read, _local_write) = local_stream.split();
        
        tokio::task::spawn_blocking(move || {
            // 在阻塞任务中处理SSH通道
            let session_guard = futures::executor::block_on(ssh_session.lock());
            
            match session_guard.channel_direct_tcpip(&remote_host, remote_port, None) {
                Ok(mut channel) => {
                    info!("建立SSH通道到 {}:{}", remote_host, remote_port);
                    
                    // 这里需要实现数据转发逻辑
                    // 由于复杂性，这里提供基本框架
                    let _ = channel.close();
                }
                Err(e) => {
                    error!("创建SSH通道失败: {}", e);
                }
            }
        }).await.map_err(|e| anyhow!("SSH通道任务失败: {}", e))?;
        
        Ok(())
    }

    /// 启动远程端口转发
    async fn start_remote_forward(
        &mut self,
        remote_port: u16,
        local_host: String,
        local_port: u16,
        ssh_session: Arc<Mutex<Session>>,
    ) -> Result<()> {
        info!("远程端口转发: {} -> {}:{}", remote_port, local_host, local_port);
        
        // 远程端口转发需要在SSH服务器上监听端口
        // 这通常需要服务器配置允许远程转发
        tokio::spawn(async move {
            let session_guard = ssh_session.lock().await;
            
            // 在阻塞任务中创建监听器
            let listener_result = tokio::task::spawn_blocking({
                let _session = session_guard.clone();
                move || {
                    // 注意：这里需要使用同步方式访问session
                    // 实际实现中可能需要不同的方法
                    info!("尝试创建远程端口转发监听器");
                    Ok::<(u16, String, u16), Box<dyn std::error::Error + Send + Sync>>((remote_port, local_host.to_string(), local_port))
                }
            });
            
            tokio::spawn(async move {
                match listener_result.await {
                    Ok(Ok((remote_port, _local_host, _local_port))) => {
                        info!("远程端口转发配置成功，远程端口: {}", remote_port);
                        
                        // 这里应该实现实际的远程端口转发逻辑
                        // 由于SSH2库的限制，完整实现需要更复杂的处理
                        loop {
                            // 模拟连接处理
                            tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
                            
                            // 在实际实现中，这里会处理来自SSH服务器的连接
                            // 并转发到本地服务
                            debug!("远程端口转发运行中...");
                        }
                    }
                    Ok(Err(e)) => {
                        error!("创建远程端口转发失败: {:?}", e);
                    }
                    Err(e) => {
                        error!("任务执行失败: {}", e);
                    }
                }
            });
        });
        
        Ok(())
    }

    /// 启动动态端口转发 (SOCKS代理)
    async fn start_dynamic_forward(
        &mut self,
        local_port: u16,
        ssh_session: Arc<Mutex<Session>>,
    ) -> Result<()> {
        let bind_addr = format!("127.0.0.1:{}", local_port);
        let listener = TokioTcpListener::bind(&bind_addr).await
            .map_err(|e| anyhow!("绑定SOCKS代理端口失败 {}: {}", local_port, e))?;
        
        info!("SOCKS代理监听: {}", bind_addr);
        
        // 先存储 listener
        self.listener = Some(listener);
        
        // 获取 listener 的引用用于 spawn
        if let Some(ref listener) = self.listener {
            // 克隆必要的数据用于 spawn
            let listener_addr = listener.local_addr().map_err(|e| anyhow!("获取监听地址失败: {}", e))?;
            
            // 重新绑定一个新的 listener 用于 spawn
            let spawn_listener = TokioTcpListener::bind(listener_addr).await
                .map_err(|e| anyhow!("重新绑定端口失败: {}", e))?;
            
            // 启动SOCKS代理服务
            tokio::spawn(async move {
                loop {
                    match spawn_listener.accept().await {
                        Ok((stream, addr)) => {
                            debug!("接受SOCKS连接: {}", addr);
                            
                            let ssh_session = Arc::clone(&ssh_session);
                            
                            tokio::spawn(async move {
                                if let Err(e) = Self::handle_socks_connection(stream, ssh_session).await {
                                    error!("SOCKS连接处理失败: {}", e);
                                }
                            });
                        }
                        Err(e) => {
                            error!("接受SOCKS连接失败: {}", e);
                            break;
                        }
                    }
                }
            });
        }
        
        Ok(())
    }

    /// 处理SOCKS连接
    async fn handle_socks_connection(
        mut stream: TokioTcpStream,
        _ssh_session: Arc<Mutex<Session>>,
    ) -> Result<()> {
        // 实现SOCKS5协议
        // 这里提供基本框架，完整实现需要处理SOCKS5握手和请求
        
        let mut buffer = [0u8; 1024];
        let n = stream.read(&mut buffer).await
            .map_err(|e| anyhow!("读取SOCKS请求失败: {}", e))?;
        
        if n < 3 {
            return Err(anyhow!("无效的SOCKS请求"));
        }
        
        // 检查SOCKS版本
        if buffer[0] != 0x05 {
            return Err(anyhow!("不支持的SOCKS版本: {}", buffer[0]));
        }
        
        // 发送认证响应（无需认证）
        stream.write_all(&[0x05, 0x00]).await
            .map_err(|e| anyhow!("发送SOCKS认证响应失败: {}", e))?;
        
        // 读取连接请求
        let n = stream.read(&mut buffer).await
            .map_err(|e| anyhow!("读取SOCKS连接请求失败: {}", e))?;
        
        if n < 10 {
            return Err(anyhow!("无效的SOCKS连接请求"));
        }
        
        // 解析目标地址和端口
        // 这里需要完整的SOCKS5协议实现
        
        info!("SOCKS连接处理完成");
        Ok(())
    }

    /// 停止隧道
    pub async fn stop(&mut self) -> Result<()> {
        info!("停止隧道: {}", self.config.name);
        
        // 关闭监听器
        if let Some(listener) = self.listener.take() {
            drop(listener);
        }
        
        // 关闭所有活跃连接
        {
            let mut connections = self.active_connections.lock().await;
            connections.clear();
        }
        
        {
            let mut status = self.status.lock().await;
            *status = TunnelStatus::Stopped;
        }
        
        info!("隧道已停止: {}", self.config.name);
        Ok(())
    }

    /// 获取隧道状态
    pub async fn get_status(&self) -> TunnelStatus {
        let status = self.status.lock().await;
        status.clone()
    }

    /// 获取活跃连接数
    pub async fn get_active_connections_count(&self) -> usize {
        let connections = self.active_connections.lock().await;
        connections.len()
    }
}

/// 隧道管理器
pub struct TunnelManager {
    tunnels: Arc<Mutex<HashMap<String, Tunnel>>>,
}

impl TunnelManager {
    pub fn new() -> Self {
        Self {
            tunnels: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    /// 创建隧道
    pub async fn create_tunnel(&self, config: TunnelConfig) -> Result<String> {
        let tunnel_id = config.id.clone();
        let tunnel = Tunnel::new(config);
        
        {
            let mut tunnels = self.tunnels.lock().await;
            tunnels.insert(tunnel_id.clone(), tunnel);
        }
        
        info!("创建隧道成功: {}", tunnel_id);
        Ok(tunnel_id)
    }

    /// 启动隧道
    pub async fn start_tunnel(
        &self,
        tunnel_id: &str,
        ssh_session: Arc<Mutex<Session>>,
    ) -> Result<()> {
        let mut tunnels = self.tunnels.lock().await;
        let tunnel = tunnels.get_mut(tunnel_id)
            .ok_or_else(|| anyhow!("隧道不存在: {}", tunnel_id))?;
        
        tunnel.start(ssh_session).await
    }

    /// 停止隧道
    pub async fn stop_tunnel(&self, tunnel_id: &str) -> Result<()> {
        let mut tunnels = self.tunnels.lock().await;
        let tunnel = tunnels.get_mut(tunnel_id)
            .ok_or_else(|| anyhow!("隧道不存在: {}", tunnel_id))?;
        
        tunnel.stop().await
    }

    /// 删除隧道
    pub async fn remove_tunnel(&self, tunnel_id: &str) -> Result<()> {
        let mut tunnels = self.tunnels.lock().await;
        if let Some(mut tunnel) = tunnels.remove(tunnel_id) {
            tunnel.stop().await?;
        }
        
        info!("删除隧道: {}", tunnel_id);
        Ok(())
    }

    /// 获取所有隧道状态
    pub async fn get_all_tunnels(&self) -> Vec<(String, TunnelConfig, TunnelStatus)> {
        let tunnels = self.tunnels.lock().await;
        let mut result = Vec::new();
        
        for (id, tunnel) in tunnels.iter() {
            let status = tunnel.get_status().await;
            result.push((id.clone(), tunnel.config.clone(), status));
        }
        
        result
    }
}

impl Default for TunnelManager {
    fn default() -> Self {
        Self::new()
    }
}

/// 便捷函数：创建隧道
pub async fn create_tunnel(
    connection_id: String,
    local_port: u16,
    remote_host: String,
    remote_port: u16,
    tunnel_type: String,
) -> Result<String> {
    let tunnel_type = match tunnel_type.as_str() {
        "local" => TunnelType::Local {
            local_port,
            remote_host,
            remote_port,
        },
        "remote" => TunnelType::Remote {
            remote_port,
            local_host: "127.0.0.1".to_string(),
            local_port,
        },
        "dynamic" => TunnelType::Dynamic { local_port },
        _ => return Err(anyhow!("不支持的隧道类型: {}", tunnel_type)),
    };
    
    let config = TunnelConfig::new(
        format!("{}隧道-{}", tunnel_type_name(&tunnel_type), local_port),
        tunnel_type,
        connection_id,
    );
    
    Ok(config.id)
}

/// 获取隧道类型名称
fn tunnel_type_name(tunnel_type: &TunnelType) -> &'static str {
    match tunnel_type {
        TunnelType::Local { .. } => "本地",
        TunnelType::Remote { .. } => "远程",
        TunnelType::Dynamic { .. } => "动态",
    }
}